import cv2
import face_recognition
import numpy as np
import os
from config import FACE_DETECTION_CONFIDENCE, FACE_RECOGNITION_THRESHOLD, KNOWN_FACES_DIR

class FaceDetector:
    def __init__(self):
        self.known_faces = []
        self.known_names = []
        self.load_known_faces()

    def load_known_faces(self):
        """加载已知人脸数据"""
        if not os.path.exists(KNOWN_FACES_DIR):
            os.makedirs(KNOWN_FACES_DIR)

        for filename in os.listdir(KNOWN_FACES_DIR):
            if filename.endswith(('.jpg', '.jpeg', '.png')):
                image_path = os.path.join(KNOWN_FACES_DIR, filename)
                image = face_recognition.load_image_file(image_path)
                face_encoding = face_recognition.face_encodings(image)
                
                if face_encoding:
                    self.known_faces.append(face_encoding[0])
                    self.known_names.append(os.path.splitext(filename)[0])

    def detect(self, frame):
        """检测图像中的人脸"""
        # 缩小图像以提高处理速度
        small_frame = cv2.resize(frame, (0, 0), fx=0.25, fy=0.25)
        rgb_small_frame = cv2.cvtColor(small_frame, cv2.COLOR_BGR2RGB)

        # 检测人脸位置
        face_locations = face_recognition.face_locations(rgb_small_frame)
        face_encodings = face_recognition.face_encodings(rgb_small_frame, face_locations)

        faces = []
        for face_location, face_encoding in zip(face_locations, face_encodings):
            # 转换回原始图像大小
            top, right, bottom, left = [coord * 4 for coord in face_location]
            
            face_data = {
                'location': (top, right, bottom, left),
                'encoding': face_encoding
            }
            faces.append(face_data)

            # 在图像上标记人脸位置
            cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)

        return faces

    def is_stranger(self, face_data):
        """判断是否为陌生人"""
        if not self.known_faces:  # 如果没有已知人脸数据，所有人都视为陌生人
            return True

        # 与已知人脸比对
        matches = face_recognition.compare_faces(
            self.known_faces,
            face_data['encoding'],
            tolerance=FACE_RECOGNITION_THRESHOLD
        )

        return not any(matches)  # 如果没有匹配项，则为陌生人